#!/usr/bin/python

import re
import random
import getopt
import sys

m_option = 0        # 1 mix packet
                    # 0 
d_option = './'     # directory for Query file 
t_option = 0        # timestamp
n_option = 1000000        # how may user  will be create
c_option = "GenerateProvisionFile.conf" #config file
r_option = 50       #
a_option = 1        # if session mode is used, this is indicate how many acct-update mesage will be generate
s_option = 0        # session query file 
u_option = ""       # user name prefix change to
p_option = ""       # password prefix change to
P_option = 1        # how many part provision file will divide to

def usage():
    print("Function: It will generate two files, one is named \"prepare, create_x, modify_x, del_x\" and another is named \"QueryFile\". ")
    print("Usage: GenerateProvisionFile [options] ")
    print("  -m   0:decide by default is stop, 1:mixed , include start,update,stop, defalut is 0")
    print("  -d   The location of the files will be stored. The default location is the directory in which the program executes. defalut is in current directory")
    print("  -t   The interval of timestamp in generated provision records, can't be negative. ")
    print("       The default value is 0, means that there will be no timestamp in generated provision file.")
    print("  -n   The number of provision and query records will be generated, should be larger than the interval of timestamp.")
    print("       The default value is 1000000.")
    print("  -c   The full name of config file. If no config file is selected, it will generate simple query automatically. defalut is './GenerateProvisionFile.conf'")
    print("  -r   The ratio of Auth queries in the generated QueryFile, should be in 0..100. The default value is 100.")
    print("  -s   If we need generate the Query data based on session progress! ")
    print("              0--disable, If no set the '-s', this is the default value.")
    print("              1--enable and access-request including accounting session ID.")
    print("  -a   If -s is set to 1, this option is indicate how many acct-update will be generated in each session")
    print("  -u   The prefix of the UserName will be change to")
    print("  -p   The prefix of the Password will be change to")
    print("  -P   if -u or -p is set, this option will indicate how may part the provision file will divide to")
    print("  -h   Find help from usage page.")

    sys.exit(0)


def hand_arg(config, packet):
    """
    hand the variable in the packet string
    """
    packet = re.sub('\$(\w+)', 
                   lambda obj: obj.group(1) == 'random' 
                   and str(int(random.random() * 1000)), packet) 
    return packet

def read_config(filename):
    """
    read the configure file and return configure structure
    """
    try:
        fp = open(filename)
    except IOError:
        print("open file faild\n")
        sys.exit(0)
    
    config = {}
    
    for line in fp:
        ret = re.search(r'>([^#=\s]+)\s*[=]*\s*([^#=\s]*)', line) #split key and value pairs
        if ret:         #section start
            key, value = ret.groups();
            
            if key == 'End':
                break
            elif not value:     #section start
                cur_key = key;
                config[key] = {}
            else:
                config[key] = value
            #print(key, value)
            continue
        
        ret = re.search(r'(^[^#>=\s]+)\s*[=]*\s*([^#=\s]*)', line) #split key and value pairs
        if ret:         #value pair
            key, value = ret.groups()
            if value:
                assert(cur_key)         #cur_key must set 
                config[cur_key][key] = value
        #print(line)
    return config

def build_packet(config, type):
    """
    build the specified type packet depend on the configure information 
    """
    assert(config[type])
    pkt = ""
    for key in config[type]:
        if pkt:
            pkt = "%s,%s=%s" % (pkt, key, config[type][key])
        else:
            pkt = "%s=%s" % (key, config[type][key])
    pkt = "%s\n\n" % pkt
    #print(pkt)
    return pkt

def build_provision(config, pkt):
    temp = {}
    temp["create"] = build_create(config, pkt)
    temp["modify"] = build_modify(config, temp["create"])
    temp["del"] = build_del(config, temp["modify"])
    return temp

def build_create(config, pkt):
    u = re.search("User-Name=(.*?)[,\n]", pkt).group(1)
    p = re.search("Password=(.*?)[,\n]", pkt).group(1)
    command = "create %s GroupName=%s;Policy=%s;UserName=%s;Password=%s\n" % (config['Table'], config['GroupName'], config['Policy'], u, p)
    return command

def build_modify(config, pkt):
    u = re.search('UserName=(.*?)[;\n]', pkt).group(1)
    id = re.search('\d+', u).group(0)
    command = "modify %s %s -set UserName=%s%s -set Password=%s%s\n" % (config['Table'], u, u_option, id, p_option, id)
    #print(command)
    return command

def build_del(config, pkt):
    u = re.search('UserName=(.*)(;|\n|)', pkt).group(1)
    command = "delete %s %s\n" % (config['Table'], u)
    #print(command)
    return command

config = read_config(c_option)

_queryfile = "QueryFile"

#hand arg
print("hand arg of command line")
try:
    opts, args = getopt.getopt(sys.argv[1:], "m:p:a:u:t:n:c:r:s:d:P:h")
except getopt.GetoptError: 
    usage()
    sys.exit(0)

f = lambda o, a: re.sub(r'\-', lambda obj: '_', "%s%s%s" % (_queryfile, o, a))
for o, a in opts:
    if o == '-m':
        try:
            m_option = int(a)
        except ValueError:
            usage()
        if not m_option in (0, 1):
            usage()
    elif o == '-a':
        try:
            a_option = int(a)
        except ValueError:
            usage()
        if not a_option > 0:
            usage()
    elif o == '-P':
        try:
            P_option = int(a)
        except ValueError:
            usage()
        if P_option <= 0:
            usage()

    elif o == '-p':
        if a[len(a) - 1] != '/':
            p_option = "%s/" % a
        else:
            p_option = a
    elif o == '-u':
        u_option = a
    elif o == '-d':
        d_option = a
    elif o == '-t':
        try:
            t_option = int(a)
        except ValueError:
            usage()

        if t_option <= 0:
            usage()
    elif o == '-n':
        try:
            n_option = int(a)
        except ValueError:
            usage()
        if n_option <= 0:
            usage()
    elif o == '-c':
        c_option = a
    elif o == '-r':
        try:
            r_option = int(a)
        except ValueError:
            usage()
        if r_option <= 0:
            usage()
    elif o == '-s':
        try:
            s_option = int(a)
        except ValueError:
            usage()
        if not s_option in (1, 1):
            usage()
    else:
        usage()
        sys.exit(0)
    if not o in ('-d', '-u', '-p'):
        _queryfile = f(o, a)

print("do some verification.....")
#some verification
if not p_option:
    p_option = u_option

if s_option > 0 or m_option > 0:        #if session mode or mix mode is used, count of auth should equel to acct
    print("Warming : if -s is used or  -m > 0, -r should be set to 100\n")
    r_option = 100

try:
    id_start = int(config['BaseUserNameNumber'])
except:
    id_start = 0 #
    
try:
    pap_radio = float(config['PAPAuthQeryRatio']) / 100
except:
    pap_radio = 0.5

n_pap = 0
n_chap = 0

_query = []
_provision = []

print("Enter the main loop ..... ")
#generate loop    
for id in range(id_start, id_start + n_option):
    if not (id - id_start and (n_chap + n_pap) / float(id - id_start) > r_option / 100.0):      #if auth / acct_stop
        if id - id_start and n_pap / float(n_chap + n_pap)  > pap_radio:            #auth_chap / auth_pap
            packet = build_packet(config, 'Auth-CHAP')
            n_chap += 1
        else:
            packet = build_packet(config, 'Auth-PAP')
            n_pap += 1
 
        if packet:
            packet = hand_arg(config, packet)
            packet = re.sub(r'\$(@)', lambda obj: str(id), packet)
            _query.append(packet)
            #print(packet)
 
        prov = build_provision(config, packet)
        if prov:
            _provision.append(prov)
     
        if s_option == 1:
            assert(r_option == 100)
            def f(type):
                packet = build_packet(config, type)
                packet = hand_arg(config, packet)
                packet = re.sub(r'\$(@)', lambda obj: str(id), packet)
                _query.append(packet)
                #print('a', packet)
                 
            f('Acct-Start')
            for ii in range(0, a_option):
                f('Acct-Update')
            f('Acct-Stop')
             
    elif m_option == 0:
        packet = build_packet(config, 'Acct-Stop')
        packet = hand_arg(config, packet)
        packet = re.sub(r'\$(@)', lambda obj: str(id), packet)
        _query.append(packet)
        #print(packet)
         
 
if m_option == 1 and s_option == 0:       #mix mode
    def f(type):
        for id in range(id_start, id_start + n_option):
            packet = build_packet(config, type)
            if packet:
                packet = hand_arg(config, packet)
                packet = re.sub(r'\$(@)', lambda obj: str(id), packet)
                _query.append(packet)
                #print(packet)
                                     
    f("Acct-Start")
    f("Acct-Update")
    f("Acct-Stop")

print("Write query to generate file....")
#write query to query file
fp = open("%s%s" % (d_option, _queryfile), "w")
for q in _query:
    fp.write(q)
fp.close()

#write provision file
if len(u_option) == 0:
    sys.exit(0)
    pass
if P_option > n_option:
    print("-P > -n")
    sys.exit(0)

part_size = int(n_option / P_option)
fp = open("%sprepare" % d_option, "w")
for i in range(0, int(n_option / part_size)):
    command = "create %s -set Name=%s\n" % (re.sub(r'\$@', lambda obj: "", config['Policy']), config['Policy'])
    command = re.sub(r'\$@', lambda obj: str(i), command)
    fp.write(command)

for i in range(0, int(n_option / part_size)):
    command = "create %s -set Name=%s;Policy=%s\n" % (re.sub(r'\$@', lambda obj: "", config['GroupName']), config['GroupName'], config['Policy'])
    command = re.sub(r'\$@', lambda obj: str(i), command)
    fp.write(command)
fp.close()

f = lambda tr, id: re.sub(r'\$@', lambda obj: str(id), tr)
i = 0
for prov in _provision:
    part = 0
    if _provision.index(prov) % part_size == 0:
        try:
            fp_create.close()
            fp_modify.close()
            fp_del.close()
        except:
            pass
        part = int(_provision.index(prov) / part_size) 
        fp_create = open("%screate_%d" % (d_option, part), "w")
        fp_create.write("timestamp\n")
        i = 0 #to count how many record writed in current file
        fp_modify = open("%smodify_%d" % (d_option, part), "w")
        fp_del = open("%sdel_%d" % (d_option, part), "w")

    
    fp_create.write(f(prov['create'], part))
    i += 1
    if t_option and i % t_option == 0:
        fp_create.write("timestamp\n")
    fp_modify.write(f(prov['modify'], part))
    fp_del.write(f(prov['del'], part))

try:
    fp_create.close()
    fp_modify.close()
    fp_del.close()
except:
    pass


        
